Skip to content

Extend On-Device Sampling Support to more Causal Language Models #553

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 39 commits into from
Nov 1, 2025
Merged

Extend On-Device Sampling Support to more Causal Language Models #553

merged 39 commits into from
Nov 1, 2025
+9 −1

Conversation

@quic-sanising
Copy link
Contributor

@quic-sanising quic-sanising commented Sep 4, 2025 •
edited
Loading

📢 Expanded On-Device Sampling Support in QEfficient

Excited to share that On-Device Sampling—previously available only for LlamaForCausalLM—is now supported across a broader set of architectures! This enhancement brings faster, more efficient inference directly to the QAIC device.

✅ Newly Supported Architectures:

  1. FalconForCausalLM
  2. GemmaForCausalLM
  3. GPT2LMHeadModel
  4. GPTJForCausalLM
  5. GraniteForCausalLM
  6. GraniteMoeForCausalLM
  7. LlamaForCausalLM (existing)
  8. MptForCausalLM
  9. Phi3ForCausalLM
  10. Qwen2ForCausalLM

⚠️ Architectures Still Pending Support:

  1. GPTBigCodeForCausalLM
  2. InternVLChatModel
  3. MistralForCausalLM
  4. MixtralForCausalLM
  5. LlamaSwiftKVForCausalLM
  6. Grok1ModelForCausalLM

We’re actively working to extend support to these models. Contributions, feedback, and testing from the community are always welcome to help accelerate this effort!

quic-sanising and others added 30 commits June 18, 2025 13:38
@quic-sanising
Add sampler transform test
8417d8f 
Signed-off-by: quic-sanising <[email protected]>
Merge branch 'main' into ods-unit-tests
27d8dd5
Add example script
067f9b5 
Signed-off-by: sanising <[email protected]>
Update docs
931860f 
Signed-off-by: sanising <[email protected]>
Enable On Device Sampling for _continuous_batching_execution()
79b6c95 
Signed-off-by: sanising <[email protected]>
Disable On Device Sampling for _regular_model_execution()
75eac30 
Signed-off-by: sanising <[email protected]>
Use same sampling parameters for each sequence in a batch
eb6e2eb 
Signed-off-by: sanising <[email protected]>
Enable On Device Sampling for _regular_model_execution()
48b35e3 
Signed-off-by: sanising <[email protected]>
Add test for greedy sampling
c83a631 
Signed-off-by: sanising <[email protected]>
Add test for random sampling
f698a24 
Signed-off-by: sanising <[email protected]>
Remove else block
7b34a07 
Signed-off-by: sanising <[email protected]>
Merge branch 'main' into ods-unit-tests
5fa7269 
Signed-off-by: sanising <[email protected]>
Reformat code
0ee201a 
Signed-off-by: sanising <[email protected]>
@quic-sanising
Merge branch 'quic:main' into ods-unit-tests
c074768
Move sampling operations, inputs, and validation functions to utils
115505e 
Signed-off-by: sanising <[email protected]>
Change model to TinyLlama
3ac7503 
Signed-off-by: sanising <[email protected]>
Add header
02669e0 
Signed-off-by: sanising <[email protected]>
Reformat code
137cc4a 
Signed-off-by: sanising <[email protected]>
@quic-sanising
Merge branch 'quic:main' into ods-unit-tests
54a926a
Update linter
6acf446 
Signed-off-by: sanising <[email protected]>
@quic-sanising
Merge branch 'quic:main' into ods-unit-tests
6083f5b
Remove device_id
c2d7e83 
Signed-off-by: sanising <[email protected]>
Remove redundant line
1069109 
Signed-off-by: sanising <[email protected]>
@quic-sanising
Merge branch 'quic:main' into ods-unit-tests
7d67132
Merge branch 'main' into ods-unit-tests
0e3f257
Remove redundant reinitialization of output buffers
908e67e 
Signed-off-by: sanising <[email protected]>
Merge branch 'main' into ods-unit-tests
a8e55da
Add qaic_config to model hash
f3f89d3 
Signed-off-by: sanising <[email protected]>
Merge branch 'main' into ods-unit-tests
81ae15a
Change config
c485bfd 
Signed-off-by: sanising <[email protected]>
@quic-sanising quic-sanising changed the title Extend On Device Sampling Support to more Causal Language Models Extend On-Device Sampling Support to more Causal Language Models Sep 4, 2025
@quic-sanising
Copy link
Contributor Author

Depends on PR #463.

@quic-sanising quic-sanising changed the base branch from main to ods-unit-tests September 4, 2025 20:26
@quic-sanising quic-sanising changed the base branch from ods-unit-tests to main September 4, 2025 20:26
@quic-sanising quic-sanising marked this pull request as ready for review September 18, 2025 18:54
@quic-hemagnih
Copy link
Contributor

Have we tested these models for On-Device Sampling? Have you added test cases in the CI for these models?

@quic-hemagnih
Copy link
Contributor

Also please rebase it, post your testing confirmation we can go ahead and merge this PR

quic-hemagnih
quic-hemagnih requested changes Oct 30, 2025
View reviewed changes
Copy link
Contributor

@quic-hemagnih quic-hemagnih left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please confirm -

  1. have you tested the newly added models for On Device Sampling
  2. Please rebase
  3. Can you add few models in the CI

@quic-sanising
Copy link
Contributor Author

  1. Yes, I have tested it locally and the feature works for each of the 10 architectures mentioned above.
  2. Rebase done.
  3. Afaik for the CI, you guys prefer lightweight models. And maybe for each new architecture, we can add one model config that can be tested (similar to Tinyllama model configs for LlamaForCausalLM architecture in the current CI). If this sounds good, can you please provide me a list of models that you want to be tested? Also, we need to keep in mind that for each new model config, we would need to add 2 sets of ground truth: one for greedy sampling and one for random sampling. Let me know how you want to proceed.

@quic-hemagnih quic-hemagnih merged commit 35d8fd8 into quic:main Nov 1, 2025
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@quic-hemagnih quic-hemagnih quic-hemagnih approved these changes

@quic-rishinr quic-rishinr Awaiting requested review from quic-rishinr quic-rishinr is a code owner

@ochougul ochougul Awaiting requested review from ochougul ochougul is a code owner

@quic-amitraj quic-amitraj Awaiting requested review from quic-amitraj quic-amitraj is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@quic-sanising @quic-hemagnih @quic-dhirajku